import torch
from torch import nn
from torchsummary import summary
import torch.nn.functional as F
class PointNetwork(nn.Module):
    def __init__(self):
        super(PointNetwork, self).__init__()
        self.fc1 = nn.Linear(210, 64)  # 输入层，将 21x3 的数据展平
        self.fc2 = nn.Linear(64, 32)  # 隐藏层
        self.fc3 = nn.Linear(32, 10)  # 输出层，假设有 10 个类别

    def forward(self, x):
        x = x.view(-1, 210)  # 将输入数据展平
        x = F.relu(self.fc1(x))  # 使用 ReLU 激活函数
        x = F.relu(self.fc2(x))
        x = self.fc3(x)  # 输出层，不需要激活函数
        return x
def load_pth(point_model, model_path):
    
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = torch.load(model_path, map_location=device)
    model.to(device)
    model.eval()  # 设置模型为评估模式
    return model

if __name__ == '__main__':
    point_model = PointNetwork()
    model_path = 'models/points.pth'  # 替换为你的模型路径
    model = load_pth(point_model, model_path)
    summary(model, (21, 4))
    print("Model loaded successfully!")